package com.wistech.search.solr4.cloud;

import java.io.IOException;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Random;

import javax.sql.DataSource;

import org.apache.solr.client.solrj.SolrServer;
import org.apache.solr.client.solrj.impl.CloudSolrServer;
import org.apache.solr.client.solrj.SolrServerException;
import org.apache.solr.common.SolrInputDocument;
import org.junit.BeforeClass;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.test.context.ContextConfiguration;
import org.springframework.test.context.junit4.SpringJUnit4ClassRunner;

import com.wistech.search.solr4.ServerFactoryImpl;
import com.wistech.search.solr4.SolrHelper;
import com.wistech.utils.DBUtils;

@ContextConfiguration(locations = { "/applicationContext-database.xml" })
@RunWith(SpringJUnit4ClassRunner.class)
public class CloudProductIndex {
	private static final Logger log = LoggerFactory
			.getLogger(CloudProductIndex.class);
	private static final String pSQL = "SELECT id,cat_id AS categoryId,company_id AS companyId ,name,publish_date AS publishDate FROM product WHERE id>? LIMIT ?";
	private static SolrServer solrServer;
	private static String host = "10.36.64.107:9983";
	private static SolrHelper helper;;
	private static final int fetchSize = 500; // 每次抓取记录条数
	@Autowired
	private DataSource dataSource;

	@BeforeClass
	public static void beforeClass() {
		ServerFactoryImpl factory = new ServerFactoryImpl();
		factory.setServerUrl(host);
		factory.setClouded(true);
		factory.buildServer();

		solrServer = factory.getSolrServer();
		((CloudSolrServer) solrServer).setDefaultCollection("collection1");
		helper = new SolrHelper(solrServer);

	}

	@Test
	public void index() {
		helper.clean();
		log.info("Loading product data from database...");
		long totalCount = getTotalCount();
		if (totalCount <= 0) {
			log.info("have no product data.");
			return;
		}
		Connection conn = null;
		PreparedStatement stmt = null;
		ResultSet rs = null;
		try {
			int pageCount = calcPageCount(totalCount, fetchSize);
			conn = dataSource.getConnection();
			stmt = conn.prepareStatement(pSQL);
			int maxId = 0;
			int id = 0;
			for (int pageNo = 1; pageNo < pageCount + 1; pageNo++) {
				log.info("共 " + totalCount + " 条记录，正在处理第 "
						+ (pageNo * fetchSize) + " 条数据...");
				long time1 = System.currentTimeMillis();
				stmt.clearParameters();
				stmt.setLong(1, maxId);
				stmt.setLong(2, fetchSize);
				rs = stmt.executeQuery();
				long time2 = System.currentTimeMillis();
				Collection<SolrInputDocument> docs = new ArrayList<SolrInputDocument>();

				while (rs.next()) {
					maxId = rs.getInt("id");
					int repeat = 10;
					for (int i = 0; i < repeat; i++) {
						id++;
						SolrInputDocument doc = toDoc(rs);
						doc.addField("id", String.valueOf(id));
						docs.add(doc);
					}
				}
				solrServer.add(docs);
				solrServer.commit();
				long time3 = System.currentTimeMillis();
				log.info("从数据库读数据耗时 " + (time2 - time1) + " 纳秒. 索引耗时  "
						+ (time3 - time2) + " 纳秒.");
				System.out.println();
			}
			log.info("全部完成.");
		} catch (SQLException e) {
			log.error("Failed to load product data.", e);
		} catch (IOException e) {
			log.error("Failed to load product data.", e);
		} catch (SolrServerException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		} finally {
			DBUtils.close(conn, stmt, rs);
		}
	}

	private SolrInputDocument toDoc(ResultSet rs) throws SQLException,
			IOException {
		SolrInputDocument doc = new SolrInputDocument();
		doc.addField("categoryId", rs.getLong("categoryId"));
		doc.addField("companyId", rs.getLong("companyId"));
		doc.addField("name", rs.getString("name"));
		doc.addField("publishDate", rs.getDate("publishDate"));
		doc.addField("orderNum", randomInt(1000));
		doc.addField("viewNum", randomInt(100000));
		return doc;
	}

	private long getTotalCount() {
		Connection conn = null;
		PreparedStatement stmt = null;
		ResultSet rs = null;
		try {
			String sql = "SELECT COUNT(id) AS totalCount FROM product";
			conn = dataSource.getConnection();
			stmt = conn.prepareStatement(sql);
			rs = stmt.executeQuery();
			if (rs.next()) {
				return rs.getLong("totalCount");
			} else {
				return 0;
			}
		} catch (SQLException e) {
			log.error("数据库访问异常", e);
		} finally {
			DBUtils.close(conn, stmt, rs);
		}
		return 0;
	}

	private int randomInt(int max) {
		Random rdm = new Random();
		int result = 0;
		do {
			result = rdm.nextInt(max);
		} while (result <= 0);
		return result;
	}

	private int calcPageCount(long totalCount, int pageSize) {
		if (totalCount == 0)
			return 0;
		int count = (int) (totalCount / pageSize);
		if (totalCount % pageSize > 0) {
			count++;
		}
		return count;
	}
}
